package org.example.brave.spring.web;

import brave.Span;
import brave.Tracing;
import brave.propagation.CurrentTraceContext;
import brave.propagation.TraceContext;
import org.example.brave.http.HttpClientHandler;
import org.example.brave.http.HttpClientRequest;
import org.example.brave.http.HttpClientResponse;
import org.example.brave.http.HttpTracing;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpRequest;
import org.springframework.http.client.AsyncClientHttpRequestExecution;
import org.springframework.http.client.AsyncClientHttpRequestInterceptor;
import org.springframework.http.client.ClientHttpResponse;
import org.springframework.util.concurrent.ListenableFuture;
import org.springframework.util.concurrent.ListenableFutureCallback;

import java.io.IOException;

public final class TracingAsyncClientHttpRequestInterceptor
        implements AsyncClientHttpRequestInterceptor {

    public static AsyncClientHttpRequestInterceptor create(Tracing tracing) {
        return create(HttpTracing.create(tracing));
    }

    public static AsyncClientHttpRequestInterceptor create(HttpTracing httpTracing) {
        return new TracingAsyncClientHttpRequestInterceptor(httpTracing);
    }

    final CurrentTraceContext currentTraceContext;
    final HttpClientHandler<HttpClientRequest, HttpClientResponse> handler;

    @Autowired
    TracingAsyncClientHttpRequestInterceptor(HttpTracing httpTracing) {
        currentTraceContext = httpTracing.tracing().currentTraceContext();
        handler = HttpClientHandler.create(httpTracing);
    }

    @Override public ListenableFuture<ClientHttpResponse> intercept(HttpRequest req,
                                                                    byte[] body, AsyncClientHttpRequestExecution execution) throws IOException {
        TracingClientHttpRequestInterceptor.HttpRequestWrapper request = new TracingClientHttpRequestInterceptor.HttpRequestWrapper(req);
        Span span = handler.handleSend(request);

        // avoid context sync overhead when we are the root span
        TraceContext invocationContext = span.context().parentIdAsLong() != 0
                ? currentTraceContext.get()
                : null;

        try (CurrentTraceContext.Scope ws = currentTraceContext.maybeScope(span.context())) {
            ListenableFuture<ClientHttpResponse> result = execution.executeAsync(req, body);
            result.addCallback(new TraceListenableFutureCallback(request, span, handler));
            return invocationContext != null
                    ? new TraceContextListenableFuture<>(result, currentTraceContext, invocationContext)
                    : result;
        } catch (Throwable e) {
            handler.handleReceive(new TracingClientHttpRequestInterceptor.ClientHttpResponseWrapper(request, null, e), span);
            throw e;
        }
    }

    static final class TraceListenableFutureCallback
            implements ListenableFutureCallback<ClientHttpResponse> {
        final TracingClientHttpRequestInterceptor.HttpRequestWrapper request;
        final Span span;
        final HttpClientHandler<HttpClientRequest, HttpClientResponse> handler;

        TraceListenableFutureCallback(
                TracingClientHttpRequestInterceptor.HttpRequestWrapper request, Span span,
                HttpClientHandler<HttpClientRequest, HttpClientResponse> handler) {
            this.request = request;
            this.span = span;
            this.handler = handler;
        }

        @Override public void onFailure(Throwable ex) {
            handler.handleReceive(new TracingClientHttpRequestInterceptor.ClientHttpResponseWrapper(request, null, ex), span);
        }

        @Override public void onSuccess(ClientHttpResponse response) {
            handler.handleReceive(new TracingClientHttpRequestInterceptor.ClientHttpResponseWrapper(request, response, null), span);
        }
    }
}
